#pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable
#pragma OPENCL EXTENSION cl_arm_printf : enable

__constant sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;

constant float At[24] = {
    1.0000000000f, 1.0000000000f,  1.0000000000f, 1.0000000000f,  1.0000000000f, 0.0000000000f,  // 1,      1,          1,          1,          1,          0
    0.0000000000f, 0.7071067691f, -0.7071067691f, 1.4142135382f, -1.4142135382f, 0.0000000000f,  // 0,      sqrt(2)/2  -sqrt(2)/2   sqrt(2)    -sqrt(2)     0
    0.0000000000f, 0.4999999702f,  0.4999999702f, 1.9999998808f,  1.9999998808f, 0.0000000000f,  // 0,      0.5,        0.5,        2,          2,          0
    0.0000000000f, 0.3535533547f, -0.3535533547f, 2.8284268379f, -2.8284268379f, 1.0000000000f   // 0,      sqrt(2)/4  -sqrt(2)/4   sqrt(2)*2  -sqrt(2)*2   0
};

__kernel void winograd_36_to_4x4_image2d(__write_only image2d_t output_tex2d,
                                       __read_only image2d_t input_tex2d,
                                       __read_only image2d_t at_tex2d,
                                       __global float4 *bias_buffer,
                                       int4 input_sliced_shape,     // 1*36*14400*64
                                       int4 output_sliced_shape) {  // 1*480*480*64
    int tile_id = get_global_id(0);         // 这个kernel是要算输出中第slice个C4channel中第tile_id个tile的第tile_row行的4个值
    int tile_row = get_global_id(1);        // tile_id也是input的width方向的id
    int slice = get_global_id(2);

    int in_height = input_sliced_shape.y, in_width = input_sliced_shape.z, in_slices = input_sliced_shape.w;
    int out_height = output_sliced_shape.y, out_width = output_sliced_shape.z, out_slices = output_sliced_shape.w;

    int tile_in_row = out_width / 4;
    int out_x = (tile_id % tile_in_row) * 4;
    int out_y = (tile_id / tile_in_row) * 4 + tile_row;

    if (out_x >= out_width || out_y >= out_height || slice >= out_slices) return;
    
    float at_row[6];        // 从at中读取第tile_row行
    float4 at_val0 = convert_float4(read_imagef(at_tex2d, smp_none, (int2)(tile_row*2+0, 0)));
    float4 at_val1 = convert_float4(read_imagef(at_tex2d, smp_none, (int2)(tile_row*2+1, 0)));
    at_row[0] = at_val0.x;  at_row[1] = at_val0.y;  at_row[2] = at_val0.z;
    at_row[3] = at_val0.w;  at_row[4] = at_val1.x;  at_row[5] = at_val1.y;

    // 用读取的1*6的at_row向量左乘这个6*6的tile，得到一个1*6的行向量
    float4 at_mul_tile[6] = {0};
    for (int i=0; i<6; i++) {
        
        for (int j=0; j<6; j++) {
            int height_offset = i * 6 + j;
            float4 input_val = convert_float4(read_imagef(input_tex2d, smp_none, (int2)(tile_id, slice*in_height+height_offset)));
            at_mul_tile[j] += at_row[i] * input_val;
            // height_offset ++;
        }
    }

    // 用这个1*6的行向量和6*4的A矩阵计算得到1*4的行向量，即最终一个Tile中的一行
    float4 bias = bias_buffer[slice];
    for (int i=0; i<4; i++) {
        float4 intermed_mul_a = bias;
        for (int j=0; j<6; j++) {
            intermed_mul_a += at_mul_tile[j] * At[i*6+j];
        }
        write_imagef(output_tex2d, (int2)(out_x + i, slice * out_height + out_y), intermed_mul_a);
    }
}